#include "nasm/visitor/symbolCollector.h"
#include "nasm/operand/operand.h"

namespace dspnasm {

    SymbolCollector::SymbolCollector(std::shared_ptr<SymbolTable> symbolTable, std::shared_ptr<ElfFile> elfFile) {
        m_elfFile = elfFile;
        m_symbolTable = symbolTable;
        m_segmentName = "";
        m_assemblyAddress = m_offset = 0;
        m_curSize = 32;
    }

    void SymbolCollector::saveSection() {
        // 第一次遇到Section时调用该函数不进行任何处理
        if (m_segmentName.empty()) {
            return;
        }
        // 添加段信息到段表中
        m_elfFile->addShdr(m_segmentName, m_assemblyAddress, m_offset);
        // 段按4字节对其
        m_offset += (4 - m_offset % 4) % 4;
        // 更新m_dataLen为当前新段的起始逻辑地址
        m_assemblyAddress += m_offset;

        // 将当前段的偏移重置为0
        m_offset = 0;
    }

    antlrcpp::Any SymbolCollector::visitDeclare(dspnasm::NasmParser::DeclareContext *ctx) {
        if (ctx->Section()) {
            // Section Id
            saveSection();
            m_segmentName = ctx->Id()->getText();
        } else if (ctx->Global()) {
            //  Global Id
            std::string id = ctx->Id()->getText();
            m_symbolTable->addSymbol(Symbol::getGlobal(id, m_segmentName));
        } else if (ctx->Externed()) {
            // Externed Id
            std::string id = ctx->Id()->getText();
            m_symbolTable->addSymbol(Symbol::getExterned(id, m_segmentName));
        } else if (ctx->Id()) {
            // Id idTail
            auto id = ctx->Id()->getText();
            if (ctx->idTail()->Colon()) {
                m_symbolTable->addSymbol(Symbol::getLabel(id, m_segmentName, m_offset));
            } else if (ctx->idTail()->Equ()) {
                u32 data = visit(ctx->idTail()->expr()).as<u32>();
                m_symbolTable->addSymbol(Symbol::getEqu(id, m_segmentName, data));
            }
        } else if (ctx->Lbrace()) {
            u32 num = visit(ctx->num()).as<u32>();
            if (num == 16 || num == 32) {
                m_curSize = num;
            } else {
                fatalMacro("指令尺寸只能为16或32位");
            }
        } else if (ctx->Times()) {
            // Times num size data
            u32 num = visit(ctx->num()).as<u32>();
            u32 size = 0;
            if (ctx->size()->Db()) {
                size = 1;
            } else if (ctx->size()->Dw()) {
                size = 2;
            } else if (ctx->size()->Dd()) {
                size = 4;
            } else if (ctx->size()->Dq()) {
                size = 8;
            }
            m_offset += num * size;
        } else {
            // size data
            u32 size = 0;
            if (ctx->size()->Db()) {
                size = 1;
            } else if (ctx->size()->Dw()) {
                size = 2;
            } else if (ctx->size()->Dd()) {
                size = 4;
            } else if (ctx->size()->Dq()) {
                size = 8;
            }
            auto types = ctx->data()->type();
            m_offset += size * types.size();
        }
        return 0;
    }

    antlrcpp::Any SymbolCollector::visitExpr(dspnasm::NasmParser::ExprContext *ctx) {
        // expr : item ( adds item )* ;
        auto items = ctx->item();
        auto adds = ctx->adds();
        auto size = adds.size();

        u32 result = visit(items[0]).as<u32>();
        for (u32 i = 0; i < size; ++i) {
            u32 item = visit(items[i + 1]).as<u32>();
            if (adds[i]->Add()) {
                result += item;
            } else if (adds[i]->Sub()) {
                result -= item;
            }
        }
        return result;
    }

    antlrcpp::Any SymbolCollector::visitItem(dspnasm::NasmParser::ItemContext *ctx) {
        // item : factor ( muls factor)* ;
        auto factors = ctx->factor();
        auto muls = ctx->muls();
        auto size = muls.size();

        auto result = visit(factors[0]).as<u32>();
        for (u32 i = 0; i < size; ++i) {
            u32 factor = visit(factors[i + 1]).as<u32>();
            if (muls[i]->Mul()) {
                result *= factor;
            } else if (muls[i]->Div()) {
                if (factor == 0) {
                    fatalMacro("除数不能为0!");
                }
                result /= factor;
            }
        }
        return result;
    }

    antlrcpp::Any SymbolCollector::visitFactor(dspnasm::NasmParser::FactorContext *ctx) {
        // factor : val ( shift val ) * ;
        auto vals = ctx->val();
        auto shifts = ctx->shift();
        auto size = shifts.size();

        auto result = visit(vals[0]).as<u32>();
        for (u32 i = 0; i < size; ++i) {
            u32 val = visit(vals[i + 1]).as<u32>();
            if (shifts[i]->Lshift()) {
                result <<= val;
            } else if (shifts[i]->Rshift()) {
                result >>= val;
            }
        }
        return result;
    }

    antlrcpp::Any SymbolCollector::visitVal(dspnasm::NasmParser::ValContext *ctx) {
        //val : Lparen expr Rparen
        //    | Id
        //    | num
        //    ;
        if (ctx->Lparen()) {
            return visit(ctx->expr());
        } else if (ctx->Id()) {
            auto symbol = m_symbolTable->getSymbol(ctx->Id()->getText());
            if (!symbol || !symbol->m_isEqu) {
                fatalMacro(("未定义的符号" + ctx->Id()->getText()).c_str());
            }
            return vectorToU32(symbol->m_data);
        } else {
            return visit(ctx->num());
        }

    }

    antlrcpp::Any SymbolCollector::visitNum(dspnasm::NasmParser::NumContext *ctx) {
        // num : Decimal
        //     | Octal
        //     | Hex
        //     | Binary
        //     ;
        u32 data = 0;
        if (ctx->Binary()) {
            data = std::stoi(ctx->Binary()->getText().substr(2), nullptr, 2);
        } else if (ctx->Octal()) {
            data = std::stoi(ctx->Octal()->getText(), nullptr, 8);
        } else if (ctx->Decimal()) {
            data = std::stoi(ctx->Decimal()->getText(), nullptr, 10);
        } else if (ctx->Hex()) {
            data = std::stoi(ctx->Hex()->getText(), nullptr, 16);
        }
        return data;
    }

    antlrcpp::Any SymbolCollector::visitInstruction(dspnasm::NasmParser::InstructionContext *ctx) {
        m_offset += m_curSize / 8;
        return 0;
    }

    antlrcpp::Any SymbolCollector::visitProgram(NasmParser::ProgramContext *ctx) {
        NasmBaseVisitor::visitProgram(ctx);
        saveSection();
        return 0;
    }
}